/*
编译与运行
    hipcc multiProcess.cpp -o multiProcess -lmpi -lrccl
    mpirun -np 8 -H host1:4,host2:4 -x LD_LIBRARY_PATH ./multiProcess -s 1024 -i 10 -n 4
mpi参数解释：
    mpirun 是用于启动 MPI（Message Passing Interface）并行程序的命令，其常用参数含义如下：
    -np：指定要启动的进程总数。
    -H：指定要在哪些主机上启动进程，后面跟着主机名和每个主机上要启动的进程数，格式为 host1:num1,host2:num2,...。
    -x：用于传递环境变量，这里 LD_LIBRARY_PATH 是动态链接库的搜索路径。
    -H host1:4,host2:4：分配到两个节点，host1 分配 4 个进程槽位（slots），host2分配4个进程槽位，实际分配总计8个槽位（4+4）与 -np 8 完全匹配。
*/
#include <stdio.h>
#include "hip/hip_runtime.h"
#include "rccl.h"
#include <unistd.h>
#include <stdint.h>
#include <stdlib.h>
#include <sys/time.h>
#include <iostream>
#include <iomanip>
#include <vector>
#include <cmath>

// 先屏蔽return-type警告，再包含MPI头文件
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wreturn-type"
#include "mpi.h"
#pragma clang diagnostic pop

using namespace std;

// 错误检查宏定义
#define MPI_CHECK(cmd)                                                                                                 \
    do {                                                                                                               \
        int e = cmd;                                                                                                   \
        if (e != MPI_SUCCESS) {                                                                                        \
            fprintf(stderr, "MPI error at %s:%d: %d\n", __FILE__, __LINE__, e);                                         \
            exit(EXIT_FAILURE);                                                                                        \
        }                                                                                                              \
    } while (0)

#define HIP_CHECK(cmd)                                                                                                 \
    do {                                                                                                               \
        hipError_t e = cmd;                                                                                            \
        if (e != hipSuccess) {                                                                                         \
            fprintf(stderr, "HIP error at %s:%d: %s\n", __FILE__, __LINE__, hipGetErrorString(e));                      \
            exit(EXIT_FAILURE);                                                                                        \
        }                                                                                                              \
    } while (0)

#define RCCL_CHECK(cmd)                                                                                                \
    do {                                                                                                               \
        ncclResult_t r = cmd;                                                                                          \
        if (r != ncclSuccess) {                                                                                        \
            fprintf(stderr, "RCCL error at %s:%d: %s\n", __FILE__, __LINE__, ncclGetErrorString(r));                    \
            exit(EXIT_FAILURE);                                                                                        \
        }                                                                                                              \
    } while (0)

// 计时工具类（修复return警告）
class Timer {
private:
    struct timeval start_, end_;
public:
    void start() {
        gettimeofday(&start_, nullptr);
        return;  // 显式return
    }
    void stop() {
        gettimeofday(&end_, nullptr);
        return;  // 显式return
    }
    double elapsed() const {
        return (end_.tv_sec - start_.tv_sec) * 1000.0 +
               (end_.tv_usec - start_.tv_usec) / 1000.0;
    }
};

// 主机哈希计算
static uint64_t getHostHash(const char *string) {
    uint64_t result = 5381;
    for (int c = 0; string[c] != '\0'; c++) {
        result = ((result << 5) + result) ^ string[c];
    }
    return result;
}

// 获取主机名
static void getHostName(char *hostname, int maxlen) {
    gethostname(hostname, maxlen);
    for (int i = 0; i < maxlen; i++) {
        if (hostname[i] == '.') {
            hostname[i] = '\0';
            return;
        }
    }
    return;  // 显式return
}

// GPU信息打印
static void printGPUInfo(int deviceId) {
    hipDeviceProp_t prop;
    HIP_CHECK(hipGetDeviceProperties(&prop, deviceId));
    printf("GPU Device %d: %s, Memory: %.2f GB\n",
           deviceId, prop.name, prop.totalGlobalMem / (1024.0 * 1024 * 1024));
    return;  // 显式return
}

// 验证AllReduce结果
static bool verifyResults(float *hostData, int size, int rank, int nRanks) {
    for (int i = 0; i < size; i++) {
        float expected = static_cast<float>(rank * nRanks);
        if (fabs(hostData[i] - expected) > 1e-5) {
            fprintf(stderr, "Rank %d: Result verification failed at index %d: expected %.2f, got %.2f\n",
                    rank, i, expected, hostData[i]);
            return false;
        }
    }
    return true;
}

// 打印使用说明
static void printUsage(const char *progName) {
    fprintf(stderr, "Usage: %s [options]\n", progName);
    fprintf(stderr, "Options:\n");
    fprintf(stderr, "  -s <size>    Data size in KB (default: 1024)\n");
    fprintf(stderr, "  -i <iter>    Number of iterations (default: 5)\n");
    fprintf(stderr, "  -n <streams> Number of HIP streams (default: 8)\n");
    fprintf(stderr, "  -h           Show this help message\n");
    return;  // 显式return
}

int main(int argc, char *argv[]) {
    // 默认参数
    int dataSizeKB = 1024;
    int iterations = 5;
    int nStreams = 8;

    // 解析命令行参数
    int opt;
    while ((opt = getopt(argc, argv, "s:i:n:h")) != -1) {
        switch (opt) {
            case 's':
                dataSizeKB = atoi(optarg);
                break;
            case 'i':
                iterations = atoi(optarg);
                break;
            case 'n':
                nStreams = atoi(optarg);
                break;
            case 'h':
                printUsage(argv[0]);
                return 0;
            default:
                printUsage(argv[0]);
                return -1;
        }
    }

    // 计算数据大小
    int dataSize = dataSizeKB * 1024 / sizeof(float);
    size_t dataBytes = dataSize * sizeof(float);

    // MPI初始化
    int myRank, nRanks, localRank = 0;
    MPI_CHECK(MPI_Init(&argc, &argv));
    MPI_CHECK(MPI_Comm_rank(MPI_COMM_WORLD, &myRank));
    MPI_CHECK(MPI_Comm_size(MPI_COMM_WORLD, &nRanks));

    // 获取主机信息并计算本地rank
    char hostname[1024];
    getHostName(hostname, 1024);
    uint64_t hostHash = getHostHash(hostname);
    uint64_t* hostHashes = new uint64_t[nRanks];

    MPI_CHECK(MPI_Allgather(&hostHash, sizeof(uint64_t), MPI_BYTE,
                           hostHashes, sizeof(uint64_t), MPI_BYTE,
                           MPI_COMM_WORLD));

    // 计算本地rank
    for (int p = 0; p < nRanks; p++) {
        if (p == myRank) break;
        if (hostHashes[p] == hostHash) {
            localRank++;
        }
    }

    // 只在主进程打印配置信息
    if (myRank == 0) {
        printf("Running RCCL test with %d ranks\n", nRanks);
        printf("Data size: %d KB (%d elements)\n", dataSizeKB, dataSize);
        printf("Iterations: %d\n", iterations);
        printf("Number of streams: %d\n", nStreams);
        printf("----------------------------------------\n");
    }

    // 打印进程信息
    printf("[%s] Rank %d/%d, Local rank: %d\n", hostname, myRank, nRanks, localRank);

    // 设置GPU设备
    HIP_CHECK(hipSetDevice(localRank));
    printGPUInfo(localRank);

    // 创建RCCL通信器
    ncclUniqueId id;
    ncclComm_t comm;

    if (myRank == 0) {
        RCCL_CHECK(ncclGetUniqueId(&id));
    }

    // 广播RCCL ID到所有进程
    MPI_CHECK(MPI_Bcast(&id, sizeof(id), MPI_BYTE, 0, MPI_COMM_WORLD));

    // 初始化RCCL通信器
    RCCL_CHECK(ncclCommInitRank(&comm, nRanks, id, myRank));

    // 分配设备内存
    float *sendbuff, *recvbuff;
    HIP_CHECK(hipMalloc(&sendbuff, dataBytes));
    HIP_CHECK(hipMalloc(&recvbuff, dataBytes));

    // 初始化发送缓冲区
    vector<float> hostData(dataSize, static_cast<float>(myRank));
    HIP_CHECK(hipMemcpy(sendbuff, hostData.data(), dataBytes, hipMemcpyHostToDevice));

    // 创建HIP流
    vector<hipStream_t> streams(nStreams);
    for (int i = 0; i < nStreams; i++) {
        HIP_CHECK(hipStreamCreate(&streams[i]));
    }

    // 同步所有进程，准备开始计时
    MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));
    Timer timer;
    timer.start();

    // 执行RCCL AllReduce操作
    for (int i = 0; i < iterations; i++) {
        RCCL_CHECK(ncclGroupStart());
        RCCL_CHECK(ncclAllReduce(sendbuff, recvbuff, dataSize,
                                ncclFloat, ncclSum, comm, streams[i % nStreams]));
        RCCL_CHECK(ncclGroupEnd());
    }

    // 同步所有流
    for (auto& stream : streams) {
        HIP_CHECK(hipStreamSynchronize(stream));
    }

    // 停止计时
    timer.stop();
    MPI_CHECK(MPI_Barrier(MPI_COMM_WORLD));

    // 计算性能
    double totalTime = timer.elapsed();
    double avgTime = totalTime / iterations;
    double bandwidthGB = (2.0 * nRanks * dataSizeKB * iterations) / (totalTime * 1024.0);

    // 收集结果到主机
    vector<float> result(dataSize);
    HIP_CHECK(hipMemcpy(result.data(), recvbuff, dataBytes, hipMemcpyDeviceToHost));

    // 验证结果
    bool success = verifyResults(result.data(), dataSize, myRank, nRanks);

    // 修复MPI_Allreduce参数错误：使用指针传递
    int local_success = success ? 1 : 0;
    int global_success = 0;
    MPI_CHECK(MPI_Allreduce(&local_success, &global_success, 1,
                           MPI_INT, MPI_MIN, MPI_COMM_WORLD));

    // 主进程打印性能统计
    if (myRank == 0) {
        printf("----------------------------------------\n");
        printf("Performance Results:\n");
        printf("Total time: %.2f ms\n", totalTime);
        printf("Average time per iteration: %.2f ms\n", avgTime);
        printf("Bandwidth: %.2f GB/s\n", bandwidthGB);
        printf("AllReduce %s\n", global_success ? "succeeded" : "failed");
    }

    // 资源清理
    for (auto& stream : streams) {
        HIP_CHECK(hipStreamDestroy(stream));
    }
    HIP_CHECK(hipFree(sendbuff));
    HIP_CHECK(hipFree(recvbuff));
    RCCL_CHECK(ncclCommDestroy(comm));
    delete[] hostHashes;

    // MPI最终化
    MPI_CHECK(MPI_Finalize());

    if (myRank == 0) {
        printf("[MPI Rank %d] Exiting successfully\n", myRank);
    }
    return 0;
}
